# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：parasolver.py
    Created on：2025/4/8 14:05
    Author：(input)
    Description: 椭圆函数参数求解器
"""
import numpy as np
import scipy.special as sp
from scipy.optimize import root_scalar

class ParameterSolver:
    def __init__(self, length, height, characteristic=1.0, dimensional=True):
        """
        初始化参数求解器，输入length、height，特征长度characteristic，有无量纲dimensional
        L: 长度参数
        H: 高度参数
        C：特征长度
        D: 输入是否为有量纲量
        """
        self.L = length
        self.H = height
        self.CL = characteristic
        self.D = dimensional
        self.k_value = None  # 模数 k
        self.k_prime_value = None  # 补模数 k'
        self.C_value = None  # 参数 C

    def solve_k_and_k_prime(self):
        """
        通过方程 H*K(k) - (L/2)*K(k') = 0 求解模数 k 和补模数 k'
        """

        # 定义目标方程
        def equation(k):
            k_prime = np.sqrt(1 - k ** 2)  # 补模数 k'
            K_k = sp.ellipk(k ** 2)        # 计算 K(k)
            K_k_prime = sp.ellipk(k_prime ** 2)  # 计算 K(k')

            # 方程 H*K(k) - (L/2)*K(k') = 0
            return self.H * K_k - (self.L / 2) * K_k_prime

        # 使用数值方法求解 k
        sol = root_scalar(equation, bracket=[0, 1], method='bisect')

        if sol.converged:
            self.k_value = sol.root
            self.k_prime_value = np.sqrt(1 - self.k_value ** 2)
        else:
            raise RuntimeError("求解模数 k 和补模数 k' 时未收敛")

    def verify_k_and_k_prime(self):
        """
        验证 k^2 + k'^2 是否等于 1
        """
        if self.k_value is None or self.k_prime_value is None:
            raise ValueError("请先求解 k 和 k'")

        sum_of_squares = self.k_value ** 2 + self.k_prime_value ** 2
        if not np.isclose(sum_of_squares, 1):
            raise ValueError(f"验证失败：k^2 + k'^2 = {sum_of_squares} 不等于 1")
        else:
            print("验证成功：k^2 + k'^2 接近 1")

    def extract_common_prefix(self,a, b):
        """提取两个浮点数的小数点前缀，使得结果更稳定"""
        a_str = f"{a:.18f}"  # 保留足够多小数
        b_str = f"{b:.18f}"

        # 找到小数点后一致的前缀
        common = ''
        for ac, bc in zip(a_str, b_str):
            if ac == bc:
                common += ac
            else:
                break

        # 转回 float
        return float(common)

    def compute_C(self):
        """
        通过方程 H - C*K(k') = 0 计算 C
        """
        if self.k_prime_value is None:
            raise ValueError("请先求解 k 和 k'")

        K_k_prime = sp.ellipk(self.k_prime_value ** 2)  # 计算 K(k')
        K_k = sp.ellipk(self.k_value ** 2) # K(k)

        C_a = (self.L/2) / K_k
        C_b = self.H / K_k_prime

        abs_error = abs(C_a - C_b)
        print(f"Ca-Cb的绝对误差: {abs_error}")

        # 若输入为有量纲量
        if self.D:
            a = 2*K_k * C_a * self.CL
            b = K_k_prime * C_b * self.CL
            print(f"输入为有量纲量：\n宽:{a}m\n高:{b}m")
        else: # 无量纲量
            a = 2*K_k * C_a * self.CL
            b = K_k_prime * C_b * self.CL
            print(f"输入为无量纲量：\n宽:{a}\n高:{b}")

        self.C_value = self.extract_common_prefix(C_a, C_b)
        return self.C_value

    def verify_corner_mappings(self):
        """
        验证矩形四个角点的映射关系：
        -CK(k) = -a
        CK(k) = a
        CK(k) + iCK(k') = a + ib
        -CK(k) + iCK(k') = -a + ib
        """
        if self.k_value is None or self.k_prime_value is None or self.C_value is None:
            raise ValueError("请先求解 k, k' 和 C")

        K_k = sp.ellipk(self.k_value ** 2)
        K_k_prime = sp.ellipk(self.k_prime_value ** 2)
        C = self.C_value

        # 矩形参数
        a = self.L / 2  # 宽度的一半
        b = self.H  # 高度

        # 计算四个角点的映射
        point1 = -C * K_k  # 应为 -a
        point2 = C * K_k  # 应为 a
        point3 = C * K_k + 1j * C * K_k_prime  # 应为 a + ib
        point4 = -C * K_k + 1j * C * K_k_prime  # 应为 -a + ib

        # 预期值
        expected_point1 = -a
        expected_point2 = a
        expected_point3 = a + 1j * b
        expected_point4 = -a + 1j * b

        # 验证
        print("验证矩形角点映射：")
        print(f"点 1: -CK(k) = {point1}, 预期: {-a}")
        print(f"点 2: CK(k) = {point2}, 预期: {a}")
        print(f"点 3: CK(k) + iCK(k') = {point3}, 预期: {a + 1j * b}")
        print(f"点 4: -CK(k) + iCK(k') = {point4}, 预期: {-a + 1j * b}")

        # 检查是否符合预期
        if not np.isclose(point1, expected_point1, rtol=1e-10):
            raise ValueError(f"点 1 映射错误：{point1} != {-a}")
        if not np.isclose(point2, expected_point2, rtol=1e-10):
            raise ValueError(f"点 2 映射错误：{point2} != {a}")
        if not np.isclose(point3, expected_point3, rtol=1e-10):
            raise ValueError(f"点 3 映射错误：{point3} != {a + 1j * b}")
        if not np.isclose(point4, expected_point4, rtol=1e-10):
            raise ValueError(f"点 4 映射错误：{point4} != {-a + 1j * b}")

        print("所有角点映射验证通过！")

    def solve(self):
        """
        执行求解过程，输出 C、k、k' 并进行验证
        """
        print("开始求解模数 k 和补模数 k'...")
        self.solve_k_and_k_prime()  # 求解 k 和 k'

        print(f"求解得到的模数 k: {self.k_value}")
        print(f"求解得到的补模数 k': {self.k_prime_value}")

        # 验证 k^2 + k'^2 是否等于 1
        print("验证 k 和 k'...")
        self.verify_k_and_k_prime()

        # 计算 C
        C = self.compute_C()
        print(f"计算得到的 C: {C}")

        # 验证角点映射
        print("验证角点映射...")
        self.verify_corner_mappings()

        return self.C_value, self.k_value, self.k_prime_value

# 测试代码
if __name__ == "__main__":
    solver = ParameterSolver(length=4.0, height=2.0, dimensional=True)
    print(f"输入：宽{solver.L}高{solver.H}")
    result = solver.solve()
    if result:
        C, k, k_prime = result
        print(f"最终结果：C={C}, k={k}, k'={k_prime}")

    solver1 = ParameterSolver(length=2.0, height=1.0, characteristic=2.0, dimensional=False)
    print(f"输入：宽{solver1.L}高{solver1.H}")
    result1 = solver1.solve()
    if result1:
        C1, k1, k_prime1 = result1
        print(f"最终结果：C={C1}, k={k1}, k'={k_prime1}")